from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier

DATASET_ID = "HuggingFaceH4/ultrachat_200k"
MODEL_ID = "bigscience/bloom-3b"
DATASET_SPLIT = "train_sft"
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 2048

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, device_map="auto", torch_dtype="auto"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# Define quantization recipe
recipe = [
    SmoothQuantModifier(
        smoothing_strength=0.8,
        mappings=[
            (["re:.*query_key_value"], "re:.*input_layernorm"),
            (["re:.*dense_h_to_4h"], "re:.*post_attention_layernorm"),
        ],
    ),
    GPTQModifier(
        scheme="W8A8",
        targets="Linear",
        ignore=["lm_head"],
        dampening_frac=0.003,
    ),
]

# Load and preprocess dataset
dataset = load_dataset(DATASET_ID, split=DATASET_SPLIT)
dataset = dataset.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))

def preprocess(example):
    """Formats the messages into a simple dialogue format."""
    text = "\n".join([msg["content"] for msg in example["messages"]])
    return {"text": text}

dataset = dataset.map(preprocess)

# Apply quantization
oneshot(
    model=model,
    dataset=dataset,
    recipe=recipe,
    output_dir="bloom-3b-gptq-w8a8",
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)